from typing_extensions import TypedDict
from langgraph.graph import StateGraph,START,END

from common.ali_qwen import get_model


class State(TypedDict):
    topic:str
    joke:str
    story:str
    poem:str
    combined_output:str


llm = get_model()

def llm_call_joke(state:State):
    msg = llm.invoke(f"写一个关于{state['topic']}的笑话")
    return {"joke":msg.content}

def llm_call_story(state:State):
    msg = llm.invoke(f"写一个关于{state['topic']}的故事")
    return {"story":msg.content}

def llm_call_poem(state:State):
    msg = llm.invoke(f"写一个关于{state['topic']}的诗歌")
    return {"poem": msg.content}

def aggregator(state:State):
    combined = f"这是一个关于state['topic]的故事、笑话、诗歌！\n\n"
    combined += f"故事：\n{state['story']}\n\n"
    combined += f"笑话：\n{state['joke']}\n\n"
    combined += f"诗歌：\n{state['poem']}\n\n"
    return {"combined_output":combined}

parallel_builder = StateGraph(State)

#添加节点
parallel_builder.add_node("call_joke",llm_call_joke)
parallel_builder.add_node("call_story",llm_call_story)
parallel_builder.add_node("call_poem",llm_call_poem)
parallel_builder.add_node("aggregator",aggregator)

#建立边
parallel_builder.add_edge(START,"call_joke")
parallel_builder.add_edge(START,"call_story")
parallel_builder.add_edge(START,"call_poem")

parallel_builder.add_edge("call_joke","aggregator")
parallel_builder.add_edge("call_story","aggregator")
parallel_builder.add_edge("call_poem","aggregator")

parallel_builder.add_edge("aggregator",END)

workflow = parallel_builder.compile()
state = workflow.invoke({"topic":"cats"})
print(state["combined_output"])

